import gradio as gr
from Core.comfy_client import ComfyUIClient
import json
import os
import time


class BatchScalingInterface:
    def __init__(self):
        self.client = ComfyUIClient()
        # 目标工作流：图像缩放（预训练集处理）.json
        self.workflow_path = os.path.join(
            os.path.dirname(__file__), "..", "Comfyui", "work_flow", "图像缩放（预训练集处理）.json"
        )
        self.temp_dir = os.path.join(os.path.dirname(__file__), "..", "temp")
        os.makedirs(self.temp_dir, exist_ok=True)

    def count_images_in_folder(self, folder_path):
        """统计文件夹中图片文件的数量"""
        if not folder_path or not os.path.exists(folder_path):
            return 0
        image_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"]
        return sum(
            1 for f in os.listdir(folder_path)
            if any(f.lower().endswith(ext) for ext in image_extensions)
        )

    def get_image_files(self, folder_path):
        if not folder_path or not os.path.exists(folder_path):
            return []
        image_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"]
        files = [
            os.path.join(folder_path, f)
            for f in os.listdir(folder_path)
            if any(f.lower().endswith(ext) for ext in image_extensions)
        ]
        return sorted(files)

    def load_workflow(self):
        try:
            with open(self.workflow_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(f"Error loading workflow: {e}")
            return None

    def convert_ui_workflow_to_api_format(self, ui_workflow):
        """将UI导出的工作流转换为ComfyUI API格式"""
        api_workflow = {}

        for node in ui_workflow.get("nodes", []):
            node_id = str(node["id"])
            api_workflow[node_id] = {"class_type": node["type"], "inputs": {}}

            # 输入连接
            for input_info in node.get("inputs", []):
                if input_info.get("link") is not None:
                    link_id = input_info["link"]
                    for link in ui_workflow.get("links", []):
                        if link[0] == link_id:
                            source_node_id = str(link[1])
                            source_output_index = link[2]
                            api_workflow[node_id]["inputs"][input_info["name"]] = [
                                source_node_id, source_output_index
                            ]
                            break

            # widgets_values 映射
            if "widgets_values" in node and node["widgets_values"]:
                w = node["widgets_values"]
                if node["type"] == "Load Image Batch":
                    # [mode, seed, action, index, label, path, pattern, allow_RGBA_output, filename_text_extension]
                    if len(w) >= 9:
                        api_workflow[node_id]["inputs"]["mode"] = w[0]
                        api_workflow[node_id]["inputs"]["seed"] = w[1]
                        api_workflow[node_id]["inputs"]["index"] = w[3]
                        api_workflow[node_id]["inputs"]["label"] = w[4]
                        api_workflow[node_id]["inputs"]["path"] = w[5]
                        api_workflow[node_id]["inputs"]["pattern"] = w[6]
                        api_workflow[node_id]["inputs"]["allow_RGBA_output"] = w[7]
                        api_workflow[node_id]["inputs"]["filename_text_extension"] = w[8]
                elif node["type"] == "ImageScale":
                    # [upscale_method, width, height, crop]
                    if len(w) >= 4:
                        api_workflow[node_id]["inputs"]["upscale_method"] = w[0]
                        api_workflow[node_id]["inputs"]["width"] = w[1]
                        api_workflow[node_id]["inputs"]["height"] = w[2]
                        api_workflow[node_id]["inputs"]["crop"] = w[3]
                elif node["type"] == "SaveImage":
                    if len(w) >= 1:
                        api_workflow[node_id]["inputs"]["filename_prefix"] = w[0]

        return api_workflow

    def batch_scale_images(self, folder_path, width, height, output_prefix, progress=gr.Progress()):
        """按文件数量循环执行工作流进行图像缩放"""
        try:
            # 校验输入
            if not folder_path or not os.path.exists(folder_path):
                return "错误：请提供有效的文件夹路径", []
            try:
                width = int(width)
                height = int(height)
            except Exception:
                return "错误：宽度/高度必须是数字", []
            if width < 0 or height < 0 or width > 2048 or height > 2048:
                return "错误：宽度/高度范围为0-2048", []

            files = self.get_image_files(folder_path)
            total = len(files)
            if total == 0:
                return "错误：该文件夹内没有图片", []

            # 载入并转换工作流
            ui_workflow = self.load_workflow()
            if not ui_workflow:
                return "错误：无法加载工作流文件", []
            workflow = self.convert_ui_workflow_to_api_format(ui_workflow)

            # 设置常量入参：路径、模式、输出前缀
            if "12" in workflow:
                workflow["12"]["inputs"]["path"] = folder_path.replace("\\", "/")
                # 确保模式/pattern有值
                workflow["12"]["inputs"].setdefault("mode", "incremental_image")
                workflow["12"]["inputs"].setdefault("pattern", "*")
            if "11" in workflow and output_prefix:
                workflow["11"]["inputs"]["filename_prefix"] = output_prefix

            results = []
            progress(0, desc="开始批量缩放...")

            # 循环执行：index = 0..total-1
            for i in range(total):
                # 更新 index、宽度、高度
                if "12" in workflow:
                    workflow["12"]["inputs"]["index"] = i
                if "1" in workflow:
                    workflow["1"]["inputs"]["width"] = width
                    workflow["1"]["inputs"]["height"] = height

                # 提交工作流
                try:
                    resp = self.client.post_prompt(workflow)
                    if "error" in resp:
                        return f"ComfyUI错误: {resp['error']}", []
                    prompt_id = resp["prompt_id"]

                    # 轮询等待输出
                    max_retries = 40
                    retry_interval = 2
                    for retry in range(max_retries):
                        hist = self.client.get_history(prompt_id)
                        if prompt_id in hist and "outputs" in hist[prompt_id]:
                            out = hist[prompt_id]["outputs"]
                            for node_id in out:
                                if "images" in out[node_id]:
                                    for img in out[node_id]["images"]:
                                        img_bytes = self.client.get_image(img["filename"], img["subfolder"], img["type"]) 
                                        # 保存到临时目录，命名包含索引
                                        file_path = os.path.join(self.temp_dir, f"scaled_{prompt_id}_{i}_{img['filename']}")
                                        with open(file_path, "wb") as f:
                                            f.write(img_bytes)
                                        results.append(os.path.abspath(file_path))
                            break
                        time.sleep(retry_interval)
                        progress(((i + retry / max_retries) / total), desc=f"处理中 {i+1}/{total}...")

                    progress((i + 1) / total, desc=f"完成 {i+1}/{total}")
                except Exception as e:
                    return f"处理第{i+1}张时出错: {e}", []

            return f"批量缩放完成，成功处理 {total} 张图片", results
        except Exception as e:
            return f"批量缩放失败: {e}", []


def Tab_image_scaling_batch():
    interface = BatchScalingInterface()

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 批量图片缩放设置")

            folder_path_input = gr.Textbox(
                label="图片文件夹路径",
                placeholder="请输入包含图片的文件夹路径，例如：D:\\images",
                value=""
            )

            file_count_display = gr.Textbox(label="文件夹中的图片数量", interactive=False, value="0")

            # 宽度/高度输入
            width_input = gr.Slider(label="目标宽度", minimum=0, maximum=2048, step=1, value=1024)
            height_input = gr.Slider(label="目标高度", minimum=0, maximum=2048, step=1, value=1024)

            # 输出前缀
            output_prefix_input = gr.Textbox(label="输出文件名前缀", value="scaled", placeholder="输出文件前缀")

            def update_file_count(folder_path):
                return str(interface.count_images_in_folder(folder_path))

            folder_path_input.change(fn=update_file_count, inputs=[folder_path_input], outputs=[file_count_display])

        with gr.Column(scale=2):
            batch_btn = gr.Button("开始批量缩放", variant="primary", size="lg")
            status_text = gr.Textbox(label="处理状态", interactive=False, lines=3)
            result_gallery = gr.Gallery(label="处理结果预览", columns=3, show_label=True, elem_id="batch_scale_result_gallery")

    batch_btn.click(
        fn=interface.batch_scale_images,
        inputs=[folder_path_input, width_input, height_input, output_prefix_input],
        outputs=[status_text, result_gallery]
    )

    # 输出目录预览（与其他tab保持一致）
    from config.config import get_output_path
    output_dir = get_output_path()

    def get_output_images():
        if not os.path.exists(output_dir):
            return []
        imgs = []
        for f in os.listdir(output_dir):
            if f.lower().endswith((".png", ".jpg", ".jpeg")):
                imgs.append(os.path.join(output_dir, f))
        imgs.sort(key=lambda p: os.path.getmtime(p), reverse=True)
        return imgs

    with gr.Row():
        refresh_btn = gr.Button("刷新输出预览")
    with gr.Row():
        output_gallery = gr.Gallery(label="输出文件夹图片", columns=4, show_label=True, elem_id="output_gallery_scale")

    refresh_btn.click(fn=get_output_images, outputs=[output_gallery])